package com.abing.core.utils.rsatest;

import javax.crypto.BadPaddingException;
import javax.crypto.Cipher;
import javax.crypto.IllegalBlockSizeException;
import java.nio.charset.StandardCharsets;
import java.security.*;
import java.security.spec.PKCS8EncodedKeySpec;
import java.security.spec.X509EncodedKeySpec;
import java.util.Base64;
import java.util.HashMap;
import java.util.Map;

public class RSAUtils {

    // RSA 密钥长度（2048 位）
    private static final int KEY_LENGTH = 2048;
    // 加密填充模式（PKCS#1 v1.5，适合短文本）
    private static final String CIPHER_PADDING = "RSA/ECB/PKCS1Padding";
    // 签名算法（SHA-256 哈希 + RSA 签名）
    private static final String SIGN_ALGORITHM = "SHA256withRSA";

    /**
     * 生成 RSA 密钥对（公钥 + 私钥）
     * @return 包含公钥和私钥的 Map（key: "publicKey"/"privateKey"）
     */
    public static Map<String, String> generateKeyPair() throws Exception {
        KeyPairGenerator keyPairGenerator = KeyPairGenerator.getInstance("RSA");
        keyPairGenerator.initialize(KEY_LENGTH);
        KeyPair keyPair = keyPairGenerator.generateKeyPair();

        // 公钥转 Base64
        String publicKey = Base64.getEncoder().encodeToString(
            keyPair.getPublic().getEncoded()
        );
        // 私钥转 Base64
        String privateKey = Base64.getEncoder().encodeToString(
            keyPair.getPrivate().getEncoded()
        );

        Map<String, String> keyMap = new HashMap<>();
        keyMap.put("publicKey", publicKey);
        keyMap.put("privateKey", privateKey);
        return keyMap;
    }

    /**
     * 公钥加密（明文 -> 密文）
     * @param publicKeyBase64 公钥（Base64 格式）
     * @param plainText 明文（需小于等于 245 字节，2048 位密钥）
     * @return 密文（Base64 格式）
     */
    public static String encrypt(String publicKeyBase64, String plainText) throws Exception {
        PublicKey publicKey = getPublicKeyFromBase64(publicKeyBase64);
        Cipher cipher = Cipher.getInstance(CIPHER_PADDING);
        cipher.init(Cipher.ENCRYPT_MODE, publicKey);
        byte[] encryptedBytes = cipher.doFinal(plainText.getBytes(StandardCharsets.UTF_8));
        return Base64.getEncoder().encodeToString(encryptedBytes);
    }

    /**
     * 私钥解密（密文 -> 明文）
     * @param privateKeyBase64 私钥（Base64 格式）
     * @param encryptedText 密文（Base64 格式）
     * @return 明文
     */
    public static String decrypt(String privateKeyBase64, String encryptedText) throws Exception {
        PrivateKey privateKey = getPrivateKeyFromBase64(privateKeyBase64);
        Cipher cipher = Cipher.getInstance(CIPHER_PADDING);
        cipher.init(Cipher.DECRYPT_MODE, privateKey);
        byte[] decryptedBytes = cipher.doFinal(Base64.getDecoder().decode(encryptedText));
        return new String(decryptedBytes, StandardCharsets.UTF_8);
    }

    /**
     * 私钥签名（明文 -> 签名）
     * @param privateKeyBase64 私钥（Base64 格式）
     * @param plainText 明文
     * @return 签名（Base64 格式）
     */
    public static String sign(String privateKeyBase64, String plainText) throws Exception {
        PrivateKey privateKey = getPrivateKeyFromBase64(privateKeyBase64);
        Signature signature = Signature.getInstance(SIGN_ALGORITHM);
        signature.initSign(privateKey);
        signature.update(plainText.getBytes(StandardCharsets.UTF_8));
        byte[] signBytes = signature.sign();
        return Base64.getEncoder().encodeToString(signBytes);
    }

    /**
     * 公钥验签（验证签名是否合法）
     * @param publicKeyBase64 公钥（Base64 格式）
     * @param plainText 原始明文
     * @param signedText 签名（Base64 格式）
     * @return 验签结果（true：合法 / false：非法）
     */
    public static boolean verify(String publicKeyBase64, String plainText, String signedText) throws Exception {
        PublicKey publicKey = getPublicKeyFromBase64(publicKeyBase64);
        Signature signature = Signature.getInstance(SIGN_ALGORITHM);
        signature.initVerify(publicKey);
        signature.update(plainText.getBytes(StandardCharsets.UTF_8));
        byte[] signBytes = Base64.getDecoder().decode(signedText);
        return signature.verify(signBytes);
    }

    // -------------------------- 密钥转换工具 --------------------------
    private static PublicKey getPublicKeyFromBase64(String publicKeyBase64) throws Exception {
        byte[] publicKeyBytes = Base64.getDecoder().decode(publicKeyBase64);
        X509EncodedKeySpec keySpec = new X509EncodedKeySpec(publicKeyBytes);
        KeyFactory keyFactory = KeyFactory.getInstance("RSA");
        return keyFactory.generatePublic(keySpec);
    }

    private static PrivateKey getPrivateKeyFromBase64(String privateKeyBase64) throws Exception {
        byte[] privateKeyBytes = Base64.getDecoder().decode(privateKeyBase64);
        PKCS8EncodedKeySpec keySpec = new PKCS8EncodedKeySpec(privateKeyBytes);
        KeyFactory keyFactory = KeyFactory.getInstance("RSA");
        return keyFactory.generatePrivate(keySpec);
    }

    public static void main(String[] args) throws Exception {
        generateKeyPair().forEach((key, value) -> System.out.println(key + ": " + value));
    }
}